from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.svm import SVC
from sklearn.model_selection import train_test_split, StratifiedKFold, cross_val_score, GridSearchCV
from sklearn.pipeline import make_pipeline
from sklearn.metrics import accuracy_score
from sklearn.preprocessing import StandardScaler
import numpy as np
import nltk
from nltk.corpus import stopwords
from nltk.stem import PorterStemmer
import re

# 下载停用词列表
# nltk.download('stopwords')
stop_words = set(stopwords.words('chinese'))

# 文本预处理函数
def preprocess_text(text):
    # 去除标点符号和数字
    text = re.sub(r'[^\w\s]', '', text)
    text = re.sub(r'\d+', '', text)
    # 转换为小写
    text = text.lower()
    # 去除停用词
    words = text.split()
    words = [word for word in words if word not in stop_words]
    # 词干提取
    ps = PorterStemmer()
    words = [ps.stem(word) for word in words]
    return ' '.join(words)

# 示例数据和标签
data = [
    "等离子除臭设备价格",
    "等离子除臭设备公司",
    "等离子除臭设备厂家",
    "等离子消防设备厂家",
    "等离子防洪设备厂家",
    "等离子废气处理设备多少钱",
    "废钢等离子设备价格",
    "等离子灭菌设备价格",
    "等离子废气除臭设备售价",
    # ... 其他句子
]
labels = [
    "价格",
    "公司",
    "设备厂家",
    "设备厂家",
    "设备厂家",
    "价格",
    "价格",
    "价格",
    "价格",
    # ... 其他标签
]

# 预处理数据
data = [preprocess_text(sentence) for sentence in data]

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(data, labels, test_size=0.2, random_state=42, stratify=labels)

# 构建Pipeline，包括TF-IDF向量化和SVM分类器
pipeline = make_pipeline(
    TfidfVectorizer(ngram_range=(1, 3)),
    StandardScaler(with_mean=False),
    SVC(kernel='linear', probability=True)
)

# 超参数优化
param_grid = {
    'tfidfvectorizer__ngram_range': [(1, 2), (1, 3)],
    'svc__C': [0.1, 1, 10, 100]
}

# 使用分层交叉验证
stratified_k_fold = StratifiedKFold(n_splits=3, shuffle=True, random_state=42)
grid_search = GridSearchCV(pipeline, param_grid, cv=stratified_k_fold, scoring='accuracy')
grid_search.fit(X_train, y_train)

# 最佳参数
print(f"最佳参数: {grid_search.best_params_}")

# 训练最佳模型
best_model = grid_search.best_estimator_
best_model.fit(X_train, y_train)

# 预测测试集
y_pred = best_model.predict(X_test)

# 评估模型
accuracy = accuracy_score(y_test, y_pred)
print(f"模型准确率: {accuracy}")

# 交叉验证评估模型性能
cv_scores = cross_val_score(best_model, data, labels, cv=stratified_k_fold)
print(f"交叉验证准确率: {np.mean(cv_scores)}")

# 对新数据进行标签预测
new_data = ["等离子废气处理设备多少钱", "等离子防洪设备厂家", "等离子设备厂家"]
new_data = [preprocess_text(sentence) for sentence in new_data]
new_labels = best_model.predict(new_data)
print(f"新数据的标签: {new_labels}")
